# -*- coding: utf-8 -*-
# @Time    : 2023/12/10 16:28
# @Author  : Pan
# @Software: PyCharm
# @Project : AnomalyDetection
# @FileName: FreatureDataset
import os
import random
import paddle
import numpy as np
from paddle import io


class FeatureTrainDataset(io.Dataset):
    def __init__(self, normal_dir=None, abnormal_dir=None, format="i3d", video_lens=32):
        super(FeatureTrainDataset, self).__init__()
        self.format = format
        self.normal_dir = normal_dir
        self.abnormal_dir = abnormal_dir
        self.video_lens = video_lens

        self.abnormal = os.listdir(self.abnormal_dir)
        self.normal = os.listdir(self.normal_dir)

    def __getitem__(self, idx):
        idx_1 = random.randint(0, len(self.normal) - 1)
        idx_2 = random.randint(0, len(self.abnormal) - 1)
        normal, abnormal = self.load_npy(self.normal_dir, self.normal[idx_1]), self.load_npy(self.abnormal_dir, self.abnormal[idx_2])

        normal = process_feat(normal, self.video_lens).transpose(1, 0, 2)
        abnormal = process_feat(abnormal, self.video_lens).transpose(1, 0, 2)
        return {
            "normal_data": paddle.to_tensor(normal, dtype="float32"),
            "abnormal_data": paddle.to_tensor(abnormal, dtype="float32"),
            "normal_label": paddle.to_tensor(0, dtype="float32"),
            "abnormal_label": paddle.to_tensor(1, dtype="float32"),
        }

    def load_npy(self, data_dir, video):
        data = np.load(os.path.join(data_dir, video))
        if self.format == "clip":
            data = data[:, None, :]
        return data

    def __len__(self):
        return min(len(self.normal), len(self.abnormal))


class FeatureInferDataset(io.Dataset):
    def __init__(self, data_dir=None, target_file=None, format="i3d"):
        super(FeatureInferDataset, self).__init__()
        self.format = format
        target_file = open(target_file, "r")
        target = target_file.read().split("\n")
        target_file.close()
        # self.data_list = [os.path.join(data_dir, os.path.basename(item)) for item in target]
        self.data_list = [os.path.join(data_dir, os.path.basename(item) + ".npy") for item in target]

    def __getitem__(self, idx):
        data = np.load(self.data_list[idx])
        if self.format == "clip":
            data = data[:, None, :]
        data = data.transpose(1, 0, 2)
        return paddle.to_tensor(data, dtype="float32"), self.data_list[idx]

    def __len__(self):
        return len(self.data_list)


def process_feat(feat, length):
    # feat => [T, 10, 2048]
    new_feat = np.zeros((length, *feat.shape[1:])).astype(np.float32)

    r = np.linspace(0, len(feat), length + 1, dtype=np.int)
    for i in range(length):
        if r[i] != r[i + 1]:
            new_feat[i, :, :] = np.mean(feat[r[i]:r[i + 1], :, :], 0)
        else:
            new_feat[i, :, :] = feat[r[i], :, :]
    return new_feat